{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c1a795a4",
   "metadata": {},
   "source": [
    "# Launch ComfyUI on Brev!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbb2d691",
   "metadata": {},
   "source": [
    "Welcome!\n",
    "\n",
    "In this notebook, we have adapted ComfyUI's default Colab notebook to run on Brev. ComfyUI is a powerful GUI for Stable Diffusion models. It uses a node-based architecture that allows users to construct very complex image/video generation workflows by connecting different nodes, each representing a different function or operation!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16808297-1050-4866-90cf-6bd145360343",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Environment Setup\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "OPTIONS = {}\n",
    "\n",
    "WORKSPACE = 'ComfyUI'\n",
    "USE_COMFYUI_MANAGER = True  #@param {type:\"boolean\"}\n",
    "UPDATE_COMFY_UI = True  #@param {type:\"boolean\"}\n",
    "INSTALL_CUSTOM_NODES_DEPENDENCIES = True  #@param {type:\"boolean\"}\n",
    "OPTIONS['UPDATE_COMFY_UI'] = UPDATE_COMFY_UI\n",
    "OPTIONS['USE_COMFYUI_MANAGER'] = USE_COMFYUI_MANAGER\n",
    "OPTIONS['INSTALL_CUSTOM_NODES_DEPENDENCIES'] = INSTALL_CUSTOM_NODES_DEPENDENCIES"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ededb20-837d-447a-b968-489fca04802d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "if ! [ -d WORKSPACE ]; then\n",
    "    echo \"Starting initial setup of ComfyUI...\"\n",
    "    git clone https://github.com/comfyanonymous/ComfyUI\n",
    "    if [ -d ComfyUI ]; then\n",
    "        echo \"Repository cloned successfully into 'ComfyUI'.\"\n",
    "        cd ComfyUI\n",
    "    else\n",
    "        echo \"Failed to clone repository.\"\n",
    "    fi\n",
    "else\n",
    "    echo \"Workspace directory already exists.\"\n",
    "fi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77c2635a-a3fe-45e8-bc63-42d70661358d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Check and handle updates for ComfyUI\n",
    "if OPTIONS['UPDATE_COMFY_UI']:\n",
    "    !echo \"Updating ComfyUI...\"\n",
    "    !git pull\n",
    "\n",
    "!echo \"Installing dependencies...\"\n",
    "!pip install xformers!=0.0.18 -r ComfyUI/requirements.txt --extra-index-url https://download.pytorch.org/whl/cu121 --extra-index-url https://download.pytorch.org/whl/cu118 --extra-index-url https://download.pytorch.org/whl/cu117"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "478acf97-9201-4b6e-b9b8-83eb82894017",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Check if we should use ComfyUI Manager\n",
    "if 'USE_COMFYUI_MANAGER' in OPTIONS and OPTIONS['USE_COMFYUI_MANAGER']:\n",
    "    # Change directory to custom_nodes if it exists\n",
    "    %cd custom_nodes\n",
    "\n",
    "    # Ensure executable permissions for various scripts\n",
    "    if not os.path.isfile(\"ComfyUI-Manager/check.sh\"):\n",
    "        !chmod 755 ComfyUI-Manager/check.sh\n",
    "    if not os.path.isfile(\"ComfyUI-Manager/scan.sh\"):\n",
    "        !chmod 755 ComfyUI-Manager/scan.sh\n",
    "    if not os.path.isfile(\"ComfyUI-Manager/node_db/dev/scan.sh\"):\n",
    "        !chmod 755 ComfyUI-Manager/node_db/dev/scan.sh\n",
    "    if not os.path.isfile(\"ComfyUI-Manager/scripts/install-comfyui-venv-linux.sh\"):\n",
    "        !chmod 755 ComfyUI-Manager/scripts/install-comfyui-venv-linux.sh\n",
    "    if not os.path.isfile(\"ComfyUI-Manager/scripts/install-comfyui-venv-win.bat\"):\n",
    "        !chmod 755 ComfyUI-Manager/scripts/install-comfyui-venv-win.bat\n",
    "\n",
    "    # Clone the ComfyUI Manager repository if it does not exist\n",
    "    if not os.path.isdir(\"ComfyUI-Manager\"):\n",
    "        print(\"-= Initial setup ComfyUI-Manager =-\")\n",
    "        !git clone https://github.com/ltdrdata/ComfyUI-Manager /home/ubuntu/verb-workspace/ComfyUI/custom_nodes\n",
    "\n",
    "    # Change directory to ComfyUI-Manager and pull the latest changes\n",
    "    %cd ComfyUI-Manager\n",
    "    !git pull\n",
    "\n",
    "    # Navigate back to the initial directory\n",
    "    %cd ..\n",
    "\n",
    "# Return to the workspace directory\n",
    "%cd $WORKSPACE\n",
    "\n",
    "# Install dependencies for custom nodes if required\n",
    "if 'INSTALL_CUSTOM_NODES_DEPENDENCIES' in OPTIONS and OPTIONS['INSTALL_CUSTOM_NODES_DEPENDENCIES']:\n",
    "    print(\"-= Install custom nodes dependencies =-\")\n",
    "    !pip install GitPython\n",
    "    !python custom_nodes/ComfyUI-Manager/cm-cli.py restore-dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b635f87a-223d-4cbc-838f-8c96bbc048c0",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "!pip install huggingface_hub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99232352-eb07-48d7-ace1-f97a9df69a00",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!wget -c https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.ckpt -P ./models/checkpoints/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2b01186-5526-43e0-9d09-63cc34d41c48",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!wget -c https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.safetensors -P ./models/vae/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f2b03f4-90fe-467f-b020-1370ab819274",
   "metadata": {},
   "outputs": [],
   "source": [
    "import subprocess\n",
    "import threading\n",
    "import time\n",
    "import socket\n",
    "\n",
    "def check_server(port):\n",
    "    while True:\n",
    "        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:\n",
    "            result = sock.connect_ex(('127.0.0.1', port))\n",
    "            if result == 0:\n",
    "                print(\"\\nServer is up and running at http://129.80.69.99:{}\".format(port))\n",
    "                break\n",
    "            time.sleep(0.5)\n",
    "\n",
    "threading.Thread(target=check_server, daemon=True, args=(8188,)).start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35d97a31-1697-4ab3-82fb-d81fa7d43963",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!python main.py --listen 0.0.0.0 --port 8188 --dont-print-server --disable-xformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1b6853f-af32-4543-80aa-9b86b275d42a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_server(port):\n",
    "    while True:\n",
    "        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:\n",
    "            result = sock.connect_ex(('127.0.0.1', port))\n",
    "            if result == 0:\n",
    "                print(\"\\nServer is up and running at http://129.80.69.99:{}\".format(port))\n",
    "                break\n",
    "            time.sleep(0.5)\n",
    "\n",
    "threading.Thread(target=check_server, daemon=True, args=(8188,)).start()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
